import sys

sys.path.append(".")

import argparse

import torch

from src.models.model import FinetuneModel


def export(checkpoint, output_path):
    """导出onnx格式的模型"""
    model = FinetuneModel()
    model.load(torch.load(checkpoint, map_location="cpu"))

    dynamic_axes = {
        "input": {2: "height", 3: "width"},
        "bbox": {2: "height", 3: "width"},
        "hm": {2: "height", 3: "width"},
        "pool": {2: "height", 3: "width"},
    }

    model.export(output_path, (1, 3, 288, 512), ["bbox", "hm", "pool"], dynamic_axes=dynamic_axes)
    return


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", type=str)
    parser.add_argument("--output_path", type=str)
    args = parser.parse_args()

    export(args.checkpoint, args.output_path)
